// Copyright 2008 Isis Innovation Limited
#include "HomographyInit.h"
#include "SmallMatrixOpts.h"
#include <utility>
#include <TooN/se3.h>
#include <TooN/SVD.h>
#include <TooN/SymEigen.h>
#include <TooN/wls_cholesky.h>
#include "MEstimator.h"

using namespace std;
bool HomographyInit::IsHomographyInlier(Matrix<3> m3Homography, HomographyMatch match)
{
	Vector<2> v2Projected = project(m3Homography * unproject(match.v2CamPlaneFirst));
	Vector<2> v2Error = match.v2CamPlaneSecond - v2Projected;
	Vector<2> v2PixelError = match.m2PixelProjectionJac * v2Error;
	double dSquaredError = v2PixelError * v2PixelError;
	return (dSquaredError < mdMaxPixelErrorSquared);
}

double HomographyInit::MLESACScore(Matrix<3> m3Homography, HomographyMatch match)
{
	Vector<2> v2Projected = project(m3Homography * unproject(match.v2CamPlaneFirst));
	Vector<2> v2Error = match.v2CamPlaneSecond - v2Projected;
	Vector<2> v2PixelError = match.m2PixelProjectionJac * v2Error;
	double dSquaredError = v2PixelError * v2PixelError;
	if(dSquaredError > mdMaxPixelErrorSquared)
		return mdMaxPixelErrorSquared;
	else 
		return dSquaredError;
}

bool HomographyInit::Compute(vector<HomographyMatch> vMatches, double dMaxPixelError, SE3 &se3SecondFromFirst)
{
	mdMaxPixelErrorSquared = dMaxPixelError * dMaxPixelError;
	mvMatches = vMatches;

	// Find best homography from minimal sets of image matches
	BestHomographyFromMatches_MLESAC();

	// Generate the inlier set, and refine the best estimate using this
	mvHomographyInliers.clear();
	for(unsigned int i=0; i<mvMatches.size(); i++)
		if(IsHomographyInlier(mm3BestHomography, mvMatches[i]))
			mvHomographyInliers.push_back(mvMatches[i]);
	for(int iteration = 0; iteration < 5; iteration++)
		RefineHomographyWithInliers();

	// Decompose the best homography into a set of possible decompositions
	DecomposeHomography();

	// At this stage should have eight decomposition options, if all went according to plan
	if(mvDecompositions.size() != 8)
		return false;

	// And choose the best one based on visibility constraints
	ChooseBestDecomposition();

	se3SecondFromFirst = mvDecompositions[0].se3SecondFromFirst;
	return true;
}

Matrix<3> HomographyInit::HomographyFromMatches(vector<HomographyMatch> vMatches)
{
	unsigned int nPoints = vMatches.size();
	assert(nPoints >= 4);
	int nRows = 2*nPoints;
	if(nRows < 9)
		nRows = 9;
	Matrix<> m2Nx9(nRows, 9);
	for(unsigned int n=0; n<nPoints; n++)
	{
		double u = vMatches[n].v2CamPlaneSecond[0];
		double v = vMatches[n].v2CamPlaneSecond[1];

		double x = vMatches[n].v2CamPlaneFirst[0];
		double y = vMatches[n].v2CamPlaneFirst[1];

		// [u v]T = H [x y]T
		m2Nx9[n*2+0][0] = x;
		m2Nx9[n*2+0][1] = y;
		m2Nx9[n*2+0][2] = 1;
		m2Nx9[n*2+0][3] = 0;
		m2Nx9[n*2+0][4] = 0;
		m2Nx9[n*2+0][5] = 0;
		m2Nx9[n*2+0][6] = -x*u;
		m2Nx9[n*2+0][7] = -y*u;
		m2Nx9[n*2+0][8] = -u;

		m2Nx9[n*2+1][0] = 0;
		m2Nx9[n*2+1][1] = 0;
		m2Nx9[n*2+1][2] = 0;
		m2Nx9[n*2+1][3] = x;
		m2Nx9[n*2+1][4] = y;
		m2Nx9[n*2+1][5] = 1;
		m2Nx9[n*2+1][6] = -x*v;
		m2Nx9[n*2+1][7] = -y*v;
		m2Nx9[n*2+1][8] = -v;
	}

	if(nRows == 9)  
		for(int i=0; i<9; i++)  // Zero the last row of the matrix, 
			m2Nx9[8][i] = 0.0;  // TooN SVD leaves out the null-space otherwise

	// The right null-space of the matrix gives the homography...
	SVD<> svdHomography(m2Nx9);
	Vector<9> vH = svdHomography.get_VT()[8];
	Matrix<3> m3Homography(vH.get_data_ptr());
	return m3Homography;
};

// Throughout the whole thing,
// SecondView = Homography * FirstView

void HomographyInit::RefineHomographyWithInliers()
{
	WLSCholesky<9> wls;
	wls.add_prior(1.0);

	vector<Matrix<2,9> > vmJacobians;
	vector<Vector<2> > vvErrors;
	vector<double> vdErrorSquared;

	for(unsigned int i=0; i<mvHomographyInliers.size(); i++)
	{
		// First, find error.
		Vector<2> v2First = mvHomographyInliers[i].v2CamPlaneFirst;
		Vector<3> v3Second = mm3BestHomography * unproject(v2First);
		Vector<2> v2Second = project(v3Second);
		Vector<2> v2Second_real = mvHomographyInliers[i].v2CamPlaneSecond;
		Vector<2> v2Error = mvHomographyInliers[i].m2PixelProjectionJac * (v2Second_real - v2Second);

		vdErrorSquared.push_back(v2Error* v2Error);
		vvErrors.push_back(v2Error);

		Matrix<2,9> m29Jacobian;
		double dDenominator = v3Second[2];

		// Jacobians wrt to the elements of the homography:
		// For x:
		m29Jacobian[0].slice<0,3>() = unproject(v2First) / dDenominator;
		Zero(m29Jacobian[0].slice<3,3>());
		double dNumerator = v3Second[0];
		m29Jacobian[0].slice<6,3>() = -unproject(v2First) * dNumerator / (dDenominator * dDenominator);
		// For y:
		Zero(m29Jacobian[1].slice<0,3>());
		m29Jacobian[1].slice<3,3>()  = unproject(v2First) / dDenominator;;
		dNumerator = v3Second[1];
		m29Jacobian[1].slice<6,3>() = -unproject(v2First) * dNumerator / (dDenominator * dDenominator);

		vmJacobians.push_back(mvHomographyInliers[i].m2PixelProjectionJac * m29Jacobian);
	}

	// Calculate robust sigma:
	vector<double> vdd = vdErrorSquared;
	double dSigmaSquared = Tukey::FindSigmaSquared(vdd);

	// Add re-weighted measurements to WLS:
	for(unsigned int i=0; i<mvHomographyInliers.size(); i++)
	{
		double dWeight = Tukey::Weight(vdErrorSquared[i], dSigmaSquared);
		wls.add_df(vvErrors[i][0], vmJacobians[i][0], dWeight);
		wls.add_df(vvErrors[i][1], vmJacobians[i][1], dWeight);
	}
	wls.compute();
	Vector<9> v9Update = wls.get_mu();
	Matrix<3> m3Update(v9Update.get_data_ptr());
	mm3BestHomography += m3Update;
}

void HomographyInit::BestHomographyFromMatches_MLESAC()
{
	// Not many matches? Don't do ransac, throw them all in a pot and see what comes out.
	if(mvMatches.size() < 10)
	{
		mm3BestHomography = HomographyFromMatches(mvMatches);
		return;
	}

	// Enough matches? Run MLESAC.
	int anIndices[4];

	Identity(mm3BestHomography);
	double dBestError = 999999999999999999.9;

	// Do 300 MLESAC trials.
	for(int nR = 0; nR < 300 ; nR++)
	{ 
		// Find set of four unique matches
		for(int i=0; i<4; i++)
		{
			bool isUnique = false;
			int n;
			while(!isUnique)
			{
				n = rand() % mvMatches.size();
				isUnique =true;
				for(int j=0; j<i && isUnique; j++)
					if(anIndices[j] == n)
						isUnique = false;
			};
			anIndices[i] = n;
		}
		vector<HomographyMatch> vMinimalMatches;
		for(int i=0; i<4; i++)
			vMinimalMatches.push_back(mvMatches[anIndices[i]]);

		// Find a homography from the minimal set..
		Matrix<3> m3Homography = HomographyFromMatches(vMinimalMatches);

		//..and sum resulting MLESAC score
		double dError = 0.0;
		for(unsigned int i=0; i<mvMatches.size(); i++)
			dError += MLESACScore(m3Homography, mvMatches[i]);

		if(dError < dBestError)
		{
			mm3BestHomography = m3Homography;
			dBestError = dError;
		}
	};
}

void HomographyInit::DecomposeHomography()
{
	mvDecompositions.clear();
	SVD<3> svd(mm3BestHomography);
	Vector<3> v3Diag = svd.get_diagonal();
	double d1 = fabs(v3Diag[0]); // The paper suggests the square of these (e.g. the evalues of AAT)
	double d2 = fabs(v3Diag[1]); // should be used, but this is wrong. c.f. Faugeras' book.
	double d3 = fabs(v3Diag[2]);

	Matrix<3> U = svd.get_U();
	Matrix<3> V = svd.get_VT().T();

	double s = M3Det(U) * M3Det(V);

	double dPrime_PM = d2;

	int nCase;
	if(d1 != d2 && d2 != d3)
		nCase = 1;
	else if( d1 == d2 && d2 == d3)
		nCase = 3;
	else
		nCase = 2;

	if(nCase != 1)
	{
		cout << "  Homographyinit: This motion case is not implemented or is degenerate. Try again. " << endl;
		return;
	}

	double x1_PM;
	double x2;
	double x3_PM;

	// All below deals with the case = 1 case.
	// Case 1 implies (d1 != d3) 
	{ // Eq. 12
		x1_PM = sqrt((d1*d1 - d2*d2) / (d1*d1 - d3*d3));
		x2    = 0;
		x3_PM = sqrt((d2*d2 - d3*d3) / (d1*d1 - d3*d3));
	};

	double e1[4] = {1.0,-1.0,1.0,-1.0};
	double e3[4] = {1.0, 1.0, -1.0,-1.0};

	Vector<3> v3np;
	HomographyDecomposition decomposition;

	// Case 1, d' > 0:
	decomposition.d = s * dPrime_PM;
	for(int signs=0; signs<4; signs++)
	{
		// Eq 13
		Identity(decomposition.m3Rp);
		double dSinTheta = (d1 - d3) * x1_PM * x3_PM * e1[signs] * e3[signs] / d2;
		double dCosTheta = (d1 * x3_PM * x3_PM + d3 * x1_PM * x1_PM) / d2;
		decomposition.m3Rp[0][0] = dCosTheta;                
		decomposition.m3Rp[0][2] = -dSinTheta;
		decomposition.m3Rp[2][0] = dSinTheta;                
		decomposition.m3Rp[2][2] = dCosTheta;

		// Eq 14
		decomposition.v3Tp[0] = (d1 - d3) * x1_PM * e1[signs];
		decomposition.v3Tp[1] = 0.0;
		decomposition.v3Tp[2] = (d1 - d3) * -x3_PM * e3[signs];

		v3np[0] = x1_PM * e1[signs];
		v3np[1] = x2;
		v3np[2] = x3_PM * e3[signs];
		decomposition.v3n = V * v3np;

		mvDecompositions.push_back(decomposition);
	}
	// Case 1, d' < 0:
	decomposition.d = s * -dPrime_PM;
	for(int signs=0; signs<4; signs++)
	{ 
		// Eq 15
		Identity(decomposition.m3Rp, -1.0);
		double dSinPhi = (d1 + d3) * x1_PM * x3_PM * e1[signs] * e3[signs] / d2;
		double dCosPhi = (d3 * x1_PM * x1_PM - d1 * x3_PM * x3_PM) / d2;
		decomposition.m3Rp[0][0] = dCosPhi;                
		decomposition.m3Rp[0][2] = dSinPhi;
		decomposition.m3Rp[2][0] = dSinPhi;                
		decomposition.m3Rp[2][2] = -dCosPhi;

		// Eq 16
		decomposition.v3Tp[0] = (d1 + d3) * x1_PM * e1[signs];
		decomposition.v3Tp[1] = 0.0;
		decomposition.v3Tp[2] = (d1 + d3) * x3_PM * e3[signs];

		v3np[0] = x1_PM * e1[signs];
		v3np[1] = x2;
		v3np[2] = x3_PM * e3[signs];
		decomposition.v3n = V * v3np;

		mvDecompositions.push_back(decomposition);
	}

	// While we have the SVD results calculated here, store the decomposition R and t results as well..
	for(unsigned int i=0; i<mvDecompositions.size(); i++)
	{
		mvDecompositions[i].se3SecondFromFirst.get_rotation() = 
			s * U * mvDecompositions[i].m3Rp * V.T();
		mvDecompositions[i].se3SecondFromFirst.get_translation() = 
			U * mvDecompositions[i].v3Tp;
	}
}

bool operator<(const HomographyDecomposition lhs, const HomographyDecomposition rhs)
{
	return lhs.nScore < rhs.nScore;
}

static double SampsonusError(Vector<2> &v2Dash, const Matrix<3> &m3Essential, Vector<2> &v2)
{
	Vector<3> v3Dash = unproject(v2Dash);
	Vector<3> v3 = unproject(v2);  

	double dError = v3Dash * m3Essential * v3;

	Vector<3> fv3 = m3Essential * v3;
	Vector<3> fTv3Dash = m3Essential.T() * v3Dash;

	Vector<2> fv3Slice = fv3.slice<0,2>();
	Vector<2> fTv3DashSlice = fTv3Dash.slice<0,2>();

	return (dError * dError / (fv3Slice * fv3Slice + fTv3DashSlice * fTv3DashSlice));
}


void HomographyInit::ChooseBestDecomposition()
{
	assert(mvDecompositions.size() == 8);
	for(unsigned int i=0; i<mvDecompositions.size(); i++)
	{
		HomographyDecomposition &decom = mvDecompositions[i];
		int nPositive = 0;
		for(unsigned int m=0; m<mvHomographyInliers.size(); m++)
		{
			Vector<2> &v2 = mvHomographyInliers[m].v2CamPlaneFirst;
			double dVisibilityTest = (mm3BestHomography[2][0] * v2[0] + mm3BestHomography[2][1] * v2[1] + mm3BestHomography[2][2]) / decom.d;
			if(dVisibilityTest > 0.0)
				nPositive++;
		};
		decom.nScore = -nPositive;
	}

	sort(mvDecompositions.begin(), mvDecompositions.end());
	mvDecompositions.resize(4);

	for(unsigned int i=0; i<mvDecompositions.size(); i++)
	{
		HomographyDecomposition &decom = mvDecompositions[i];
		int nPositive = 0;
		for(unsigned int m=0; m<mvHomographyInliers.size(); m++)
		{
			Vector<3> v3 = unproject(mvHomographyInliers[m].v2CamPlaneFirst);
			double dVisibilityTest = v3 * decom.v3n / decom.d;
			if(dVisibilityTest > 0.0)
				nPositive++;
		};
		decom.nScore = -nPositive;
	}

	sort(mvDecompositions.begin(), mvDecompositions.end());
	mvDecompositions.resize(2);

	// According to Faugeras and Lustman, ambiguity exists if the two scores are equal
	// but in practive, better to look at the ratio!
	double dRatio = (double) mvDecompositions[1].nScore / (double) mvDecompositions[0].nScore;

	if(dRatio < 0.9) // no ambiguity!
		mvDecompositions.erase(mvDecompositions.begin() + 1);
	else             // two-way ambiguity. Resolve by sampsonus score of all points.
	{
		double dErrorSquaredLimit  = mdMaxPixelErrorSquared * 4;
		double adSampsonusScores[2];
		for(int i=0; i<2; i++)
		{
			SE3 se3 = mvDecompositions[i].se3SecondFromFirst;
			Matrix<3> m3Essential;
			for(int j=0; j<3; j++)
				m3Essential.T()[j] = se3.get_translation() ^ se3.get_rotation().get_matrix().T()[j];

			double dSumError = 0;
			for(unsigned int m=0; m < mvMatches.size(); m++ )
			{
				double d = SampsonusError(mvMatches[m].v2CamPlaneSecond, m3Essential, mvMatches[m].v2CamPlaneFirst);
				if(d > dErrorSquaredLimit)
					d = dErrorSquaredLimit;
				dSumError += d;
			}

			adSampsonusScores[i] = dSumError;
		}

		if(adSampsonusScores[0] <= adSampsonusScores[1])
			mvDecompositions.erase(mvDecompositions.begin() + 1);
		else
			mvDecompositions.erase(mvDecompositions.begin());
	}

}
